#!/usr/bin/env python3
# encoding: utf-8
"""
@author: 孙伟伟
@contact: jakesun2020@163.com
@application: Embeddding层将称作文本嵌入层, Embedding层产生的张量称为词嵌入张量, 它的最后一维将称作词向量
@file: sorce_code.py
@time: 2021/11/22/022 11:05
"""
# 导入必备的工具包
from numpy.lib.utils import source
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import math
# torch中变量封装函数Variable.
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F

'''无论是源文本嵌入还是目标文本嵌入，都是为了将文本中词汇的数字表示转变为向量表示, 希望在这样的高维空间捕捉词汇间的关系'''
# 1词嵌入embedding， 最后一维称为词向量
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        '''
        :param vocab:
        :param word_dim: 词嵌入的维度，即一个单词由多语料库大小，词表大小少列组成，通常是512维，将词向量的维度从vocab转换到d_model
        '''
        super(Embeddings, self).__init__()
        self.vocab = vocab
        self.d_model = d_model
        self.embed = nn.Embedding(vocab, d_model)

    def forward(self, x):
        # x表示输入文本经过词汇映射后的数字张量
        return self.embed(x) * math.sqrt(self.d_model)# 乘上维度的根号进行缩放

# embedding = nn.Embedding(10, 3)
# input = torch.LongTensor([[1,2,4,5],[4,3,2,9]])
# print(embedding(input))
'''
经过embedding之后，变成了3列，即一行表示一个单词，其中一个单词有3列表示

tensor([[[-0.0734, -0.6589,  0.5305],   1   
         [-1.1470,  1.3824, -1.4246],   2
         [ 1.3420, -0.7482,  0.1109],   4
         [-0.6119, -0.1617,  2.7381]],  5

        [[ 1.3420, -0.7482,  0.1109],   4
         [ 2.3577,  0.8046, -0.8422],   3
         [-1.1470,  1.3824, -1.4246],   2
         [-1.3497, -0.3085, -0.7871]]], grad_fn=<EmbeddingBackward0>)

'''
d_model = 512
vocab = 1000
# 输入x使用Variable进行封装，shape = (2, 4)
x = Variable(torch.LongTensor([[100,2,421,508],[491,998,1,221]]))
# x = x.data
emb = Embeddings(d_model, vocab)
embr = emb(x)
print("[[100,2,421,508],[491,998,1,221]]进行词嵌入的结果：\n", embr)
print("-"*10 + "万用的分隔符" + "-"*10 + "\n")
'''
tensor([[[  9.2356,  10.3331, -22.6440,  ...,  17.4836, -37.4644,  26.0129],    100
         [ 26.3748, -47.3185,  11.7121,  ...,  12.1132,  29.4155,  25.2254],    2
         [ 56.9969,  -9.8070,  -4.6534,  ...,  -3.9118,  28.2038,  28.9221],    421
         [-12.9290,   1.8957, -17.7828,  ...,  10.7078, -49.5591,   1.1285]],   508

        [[ 23.2327,  -6.6930, -24.7947,  ..., -22.6785,  27.0311,  23.4831],    491
         [-22.6228,   0.6236,   9.6708,  ..., -33.4133, -23.6214,   0.5783],    998
         [ 11.6907,  10.1680,   7.6301,  ...,   4.6352,  12.7305,  27.6340],    1
         [  3.5716, -28.4926, -13.3858,  ...,   5.5262,  10.3595,  14.6809]]],  221
       grad_fn=<MulBackward0>)
[torch.FloatTensor of size 2x4x512]
'''

'''因为在Transformer的编码器结构中, 并没有针对词汇位置信息的处理，因此需要在Embedding层后加入位置编码器，将词汇位置不同可能会产生不同语义的信息加入到词嵌入张量中, 以弥补位置信息的缺失'''
# 2.位置编码器
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        '''
        初始化编码器类
        :param d_model: 词嵌入的维度
        :param dropout: 置零率，表示以百分之多少的比例让神经网络中的神经元失效
        :param max_len: 句子的最大长度
        '''
        super(PositionalEncoding, self).__init__()

        # 实例化nn中预定义的Dropout层, 并将dropout传入其中, 获得对象self.dropout
        self.dropout = nn.Dropout(p=dropout)

        # 初始化一个位置编码矩阵, 它是一个0阵，矩阵的大小是max_len x d_model.
        pe = torch.zeros(max_len, d_model) # [max_len, d_model]
        # 给长度为max_len的句子位置进行编码，即绝对编码，第一个位置为0，最后一个位置的编码为max_len-1,将其转换为 max_len * 1的形状
        position = torch.arange(0, max_len).unsqueeze(1)    # [max_len, 1]

        # 位置编码矩阵算法，这个是一个相对位置，把每个位置对应的编码压缩到0-1之间，不使用绝对编码，可以减少对文本特征的影响
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))    # [1, d_model / 2]
        pe[:, 0::2] = torch.sin(position * div_term)    # position * div_term = [max_len, d_model / 2]
        pe[:, 1::2] = torch.cos(position * div_term)
        # 得到位置编码矩阵pe
        # 因为transformer的输入是word embedding + position embedding，word embedding的维度是三维
        # 所以要给位置编码矩阵增加一个维度
        pe = pe.unsqueeze(0) # [1, max_len, d_model]

        # 最后把pe位置编码矩阵注册成模型的buffer
        # buffer：对模型效果有帮助的，但是却不是模型结构中超参数或者参数，不需要随着优化步骤进行更新的增益对象.不用反向梯度运算，更加可读
        # 注册之后我们就可以在模型保存后，重加载时和模型结构与参数一同被加载.
        self.register_buffer('pe', pe)


    def forward(self, x):
        # x:  batch_size, max_len, d_model]
        """forward函数的参数是x, 表示文本序列的词嵌入表示"""
        # transformer的输入是word/position embedding的拼接，位置嵌入一行的长度为max_len=5000，很难有一个句子的有5000个单词，
        # 因此将第二个维度，max_len对应的那个维度缩小为与word embedding的长度, word embedding的第二维即句子的最大长度
        # 最后使用Variable进行封装，使其与x的样式相同，但是它是不需要进行梯度求解的，因此把requires_grad设置成false.
        x = x + Variable(self.pe[:, :x.size(1)],
                         requires_grad=False)
        # 最后使用self.dropout对象进行随机'丢弃'操作
        return self.dropout(x)


# 词嵌入维度是512维
d_model = 512

# 置0比率为0.1
dropout = 0.1

# 句子最大长度
max_len=60

pe = PositionalEncoding(d_model, dropout, max_len)
x = embr
pe_result = pe(x)

print("pe_result词嵌入矩阵+位置嵌入矩阵:", pe_result)
print(pe_result.size())
'''
词嵌入矩阵+位置嵌入矩阵: tensor([[[ -0.0000,   9.2779,   6.2912,  ...,  46.2286,   0.0000, -27.4725],
         [ -2.2440,   3.0404,   1.6418,  ...,   9.4147,  -0.0000,  53.8078],
         [-13.2460, -15.6874, -22.7718,  ...,   3.9620, -23.7117,  41.2869],
         [-54.3887,  13.7630,   1.0295,  ...,  13.0132,   9.6554,   6.8153]],

        [[-10.7300,  31.9356,  -0.0000,  ...,   0.0000,  -0.0000,   0.8109],
         [ 38.3725,  -8.9297, -58.3472,  ...,  -8.1481,   6.0751,  -0.0000],
         [ 15.0539, -25.0214,  12.7021,  ...,   2.4765,   0.0000,  19.4092],
         [-11.0289, -16.5760,  17.8044,  ...,   5.5084,  -0.0000, -22.7275]]],
       grad_fn=<MulBackward0>)
torch.Size([2, 4, 512])
'''
print("-"*10 + "万用的分隔符" + "-"*10 + "\n")

#
#
# # 创建一张15 x 5大小的画布
# plt.figure(figsize=(15, 5))
#
# # 实例化PositionalEncoding类得到pe对象, 输入参数是20和0,表示 列20维 dropout=0
# pe = PositionalEncoding(20, 0)
#
# # 然后向pe传入被Variable封装的tensor, 这样pe会直接执行forward函数,
# # 且这个tensor里的数值都是0, 被处理后相当于位置编码张量
# y = pe(Variable(torch.zeros(1, 100, 20)))
#
# # 然后定义画布的横纵坐标, 横坐标到100的长度, 纵坐标是某一个词汇中的某维特征在不同长度下对应的值
# # 因为总共有20维之多, 我们这里只查看4，5，6，7维的值.
# plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
#
# # 在画布上填写维度提示信息
# plt.legend(["dim %d"%p for p in [4,5,6,7]])
#
# plt.show()


# 3.生成掩码张量,一个方针的下三角,0被遮掩和1被遮掩可以自定义
def subsequent_mask(size):
    attn_shape = (1,size, size)

    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')# 取上三角矩阵，不加对角线

    return torch.from_numpy(1 - subsequent_mask)# 把上三角转化为下三角，此时包含对角线上的元素

# size = 5
# sm = subsequent_mask(size)
# print(sm)
# # 掩码|张量可视化
# plt.figure(figsize=(5, 5))
# plt.imshow(subsequent_mask(20)[0])
# plt.show()


'''
# 4.注意力机制
    注意力机制计算规则 Q，K，V
        假如我们有一个问题: 给出一段文本，使用一些关键词对它进行描述!
        为了方便统一正确答案，这道题可能预先已经给大家写出了一些关键词作为提示.其中这些给出的提示就可以看作是key， 
        而整个的文本信息就相当于是query，value的含义则更抽象，可以比作是你看到这段文本信息后，脑子里浮现的答案信息，
        这里我们又假设大家最开始都不是很聪明，第一次看到这段文本后脑子里基本上浮现的信息就只有提示这些信息，  给你的提示就是全部答案
            因此key与value基本是相同的，但是随着我们对这个问题的深入理解，通过我们的思考脑子里想起来的东西原来越多，
        并且能够开始对我们query也就是这段文本，提取关键信息进行表示.  这就是注意力作用的过程， 通过这个过程，
        我们最终脑子里的value发生了变化，
        根据提示key生成了query的关键词表示方法，也就是另外一种特征表示方法.
        
        刚刚我们说到key和value一般情况下默认是相同，与query是不同的，这种是我们一般的注意力输入形式，
        但有一种特殊情况，就是我们query与key和value相同，这种情况我们称为自注意力机制，就如同我们的刚刚的例子， 
        使用一般注意力机制，是使用不同于给定文本的关键词表示它. 而自注意力机制,
        需要用给定文本自身来表达自己，也就是说你需要从给定文本中抽取关键词来表述它, 相当于对文本自身的一次特征提取.
'''
def attention(query, key, value, mask=None, dropout=None):
    # query, key, value表示注意力的三个输入张量
    # mask: 掩码张量
    # dropout: 传入的Dropout实例化对象
    d_k = query.size(-1)# query的最后一维等同于词嵌入的维度
    # 根据公式求出注意力分数， query与key的转置相乘，是将key的最后两个维度转置
    #       query的最后一个维度是词嵌入的维度，要想把两个矩阵相乘，需要把key的最后一维与前面的一个维度相互换
    # 然后除以缩放因子，即维度的根号，使得内积不会太大，太大经过softmax之后就非0即1了，不够soft
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    print("query = key = value = pe_result, 将QKV输入到attention函数中")
    print("query.size:", query.size(), '/  key.transpose(-2, -1).size: ', key.transpose(-2, -1).size())
    # query.size: torch.Size([2, 4, 512]) / key.transpose(-2, -1).size:  torch.Size([2, 512, 4])

    print("scores: ", scores)
    print("scores.size: ", scores.size())
    # scores.size::  torch.Size([2, 4, 4])

    # 判断是否使用掩码张量
    if mask is not None:
        # 使用tensor的masked_fill方法, 将掩码张量和scores张量每个位置一一比较, 如果掩码张量处为0
        # 则对应的scores张量用-1e9这个值来替换
        scores = scores.masked_fill(mask == 0, -1e9)# 上三角填充为-1e9
    # 对scores的最后一维进行softmax操作，获取最终的注意力张量
    p_attn = F.softmax(scores, dim=-1)
    # 判断是否使用dropout
    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value) ,p_attn# p_atten表示query注意力表示


# Q = K = V = 位置编码器的输出(word embedding + positional embedding得到的结果)
query = key = value = pe_result
mask = Variable(torch.zeros(2, 4, 4))

attn, p_attn = attention(query, key, value, mask)

print("query = key = value:", query)
print("scores和v相乘得到的注意力矩阵attn:", attn)
print(attn.shape)
print("q*k得到的scores经过softmax后的分数p_attn:", p_attn)
print(p_attn.shape)
'''
=====================没有经过掩码的注意力矩阵，和scores得分===================================
scores和v相乘得到的注意力矩阵attn: tensor([[[  0.0000, -11.7070,  -8.0099,  ..., -55.2664,  -6.3958,  28.9768],
         [ 22.5130,  15.0678, -22.6942,  ...,  20.0529,  18.9404,  32.4554],
         [ -7.6398,  -5.0310,  42.4178,  ...,  25.6328,   0.0000, -25.2859],
         [ 24.6805,   9.4768,   8.2109,  ...,  -7.8535,  26.4834, -55.6398]],

        [[-11.0352,  20.9319,  -6.5450,  ...,  25.6816, -21.7518,   0.0000],
         [-17.7406, -33.8542, -14.4055,  ...,   0.0000, -28.0325,  -8.6560],
         [-10.2419, -48.8036,  11.7486,  ...,  -2.8072,  -1.5940, -18.5343],
         [-27.8449,   0.0000, -36.1923,  ...,   5.6491,   5.2654,  -1.0555]]],
       grad_fn=<UnsafeViewBackward0>)
torch.Size([2, 4, 512])
q*k得到的scores经过softmax后的分数p_attn: tensor([[[1., 0., 0., 0.],
         [0., 1., 0., 0.],
         [0., 0., 1., 0.],
         [0., 0., 0., 1.]],

        [[1., 0., 0., 0.],
         [0., 1., 0., 0.],
         [0., 0., 1., 0.],
         [0., 0., 0., 1.]]], grad_fn=<SoftmaxBackward0>)
torch.Size([2, 4, 4])
======================经过mask之后的注意力矩阵，和scores得分================
scores和v相乘得到的注意力矩阵attn: tensor([[[ 23.3303, -15.9498,   4.6672,  ...,  -4.1905, -13.6027,  -7.9054],
         [ 23.3303, -15.9498,   4.6672,  ...,  -4.1905, -13.6027,  -7.9054],
         [ 23.3303, -15.9498,   4.6672,  ...,  -4.1905, -13.6027,  -7.9054],
         [ 23.3303, -15.9498,   4.6672,  ...,  -4.1905, -13.6027,  -7.9054]],

        [[ -9.8163, -16.9342, -23.7833,  ...,  -0.4832,  -3.1493,  -6.5783],
         [ -9.8163, -16.9342, -23.7833,  ...,  -0.4832,  -3.1493,  -6.5783],
         [ -9.8163, -16.9342, -23.7833,  ...,  -0.4832,  -3.1493,  -6.5783],
         [ -9.8163, -16.9342, -23.7833,  ...,  -0.4832,  -3.1493,  -6.5783]]],
       grad_fn=<UnsafeViewBackward0>)
torch.Size([2, 4, 512])
q*k得到的scores经过softmax后的分数p_attn: tensor([[[0.2500, 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500, 0.2500]],

        [[0.2500, 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500, 0.2500],
         [0.2500, 0.2500, 0.2500, 0.2500]]], grad_fn=<SoftmaxBackward0>)
torch.Size([2, 4, 4])
'''



# 多头注意力机制
'''能让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，让词义拥有来自更多元的表达，实验表明可以从而提升模型效果'''
import copy# copy进行深度copy，深度copy:重新开辟一片内存空间，把所有的内容全部复制过来
# 在多头注意力机制的实现中, 用到多个结构相同的线性层.
# 我们将使用clone函数将他们一同初始化在一个网络层列表对象中. 之后的结构中也会用到该函数
def clones(module, N):
    '''用于生成相同网络层的克隆函数, module表示要克隆的目标网络层, N代表需要克隆的数量'''
    # 在函数中, 我们通过for循环对module进行N次深度拷贝, 使其每个module成为独立的层,
    # 然后将其放在nn.ModuleList类型的列表中存放.
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

# 多头注意力机制的实现
# 这种结构设计能让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，让词义拥有来自更多元的表达
class MultiHeadeAttention(nn.Module):
    def __init__(self, head, embedding_dim, dropout=0.1):
        '''
        多头：对三个变换张量QKV分别进行线性变换，变换不会改变原有张量的尺寸，因此每个变换矩阵都是方阵
             每个头开始从从词义层面分割输出的张量，也就是每个头都想获得一组QKV进行注意力机制的计算
                但是句子中的每个词的表示只获得一部分，也就是分割了最后一维的词嵌入向量
                将每个头的获得的输入送到注意力机制中，就形成了多头注意力机制
        :param head: 头数
        :param embedding_dim: 词嵌入的维度
        :param dropout: 置0率
        '''
        super(MultiHeadeAttention, self).__init__()

        assert embedding_dim % head == 0# 进行确认多头的数量需要整除嵌入的维度
        # 得到每个头获得的分隔词向量维度d_k
        self.d_k = embedding_dim // head
        self.head = head

        # 然后获得线性层对象，通过nn的Linear实例化，
        #   它的内部变换矩阵是embedding_dim * embedding_dim，然后使用clones函数克隆四个，
        # 为什么是四个呢，这是因为在多头注意力中，Q，K，V各需要一个，最后拼接的矩阵还需要一个，因此一共是四个.
        self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)

        # 初始化注意力张量
        self.attn = None

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask = None):# 多头注意力真正计算的函数
        # 判断是否使用掩码张量
        if mask is not None:
            # 将掩码张量进行维度扩充，代表多头中的第N个头
            mask = mask.unsqueeze(0)

        batch_size = query.size(0)

        query, key, value = \
            [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
                for model, x in zip(self.linears, (query, key, value))]

        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        # 经过transpose之后，要想进行view操作，必要要进行contiguous操作才可以view
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)# head * d_k正好等于了embedding_dim

        return self.linears[-1](x)

head = 8
embeding_dim = 512

query = key = value = pe_result# 自注意力机制

mask = Variable(torch.zeros(8, 4, 4))

mha = MultiHeadeAttention(head, embeding_dim)

mha_result = mha(query, key, value, mask)
print("mha_result:", mha_result)
print("mah_result.size:", mha_result.size())
'''
mha_result: tensor([[[-3.2012, -3.7650,  3.0608,  ..., -0.7377,  6.5227,  4.2534],
         [-4.2626, -4.5817,  3.4864,  ..., -0.8767,  5.9403,  5.3380],
         [-1.0174, -5.2462,  1.4869,  ..., -0.2229,  4.1295,  3.5800],
         [-3.6492, -6.6305,  1.6812,  ...,  0.0975,  6.9458,  6.2422]],
        [[ 0.6688, -1.4791, -7.1626,  ..., -4.0287, -6.6185,  4.7661],
         [ 0.4237, -1.8200, -7.9322,  ..., -5.9500, -6.5741,  4.2853],
         [ 0.2354, -0.3358, -6.2379,  ..., -7.3107, -9.7056,  4.4648],
         [-2.1943,  0.0261, -7.5511,  ..., -6.0659, -5.0598,  4.0868]]],
       grad_fn=<AddBackward0>)
mah_result.size: torch.Size([2, 4, 512])
'''

# 5.前馈全连接层
    # 增加全连接层可以增强模型的能力
class PositionwiseForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        '''
        :param d_model: 词嵌入的维度，也是线性层的输入维度和输出维度
        :param d_ff:  第一线性层的输出，第二个线性层的输入
        :param dropout:
        '''
        super(PositionwiseForward, self).__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x表示上一层的输出，多头注意力的输出
        # 首先经过一个线性层，然后经过relu激活
        # 使用dropout进行随机置零，再通过第二个线性层
        return self.w2(self.dropout(F.relu(self.w1(x))))

print(20*"=" + "万能的分隔符"+ 20*"=")

d_model = 512
d_ff = 64

x = pe_result

ff = PositionwiseForward(d_model, d_ff)
ff_result = ff(x)
print("ff_result:", ff_result)
print("ff_result.size:", ff_result.size())
'''
ff_result: tensor([[[ -5.1375,  10.4738, -11.7356,  ...,   1.4996,   1.6599,   7.5966],
         [ -2.3896,   9.0038,   7.6251,  ...,  -5.0198,  -1.5134,   7.1911],
         [  5.4251,   7.4480,  -7.5439,  ...,   5.1491,  12.7045,  -3.8947],
         [ -0.1843,  11.8629,   1.8759,  ...,   3.0128,   0.5877,  -3.8646]],
        [[ -7.5043,  10.6796,  -4.8167,  ...,   4.3128,  -3.9490, -12.6748],
         [ -7.7593,  -2.6423,   0.9214,  ...,  -7.3377,  -1.5898,  -1.1688],
         [ -6.9705,   8.7540,  -9.8427,  ...,   7.4905,   6.1496,  -0.4904],
         [  1.6040,   7.4654,  -0.4437,  ...,  -3.3194,  -4.2475,  -0.6594]]],
       grad_fn=<AddBackward0>)
ff_result.size: torch.Size([2, 4, 512])
'''

# 6.归一化层
class LayerNorm(nn.Module):
    def __init__(self, features, ep=1e-6):
        # features: 词嵌入维度
        # ep: 防止除零的操作
        super(LayerNorm, self).__init__()
        # 初始化两个参数张量，用于对结果做规范化计算
        # 用nn.parameter进行封装，代表他们也是模型中的参数
        self.a2 = nn.Parameter(torch.ones(features))
        self.b2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        # 只是对值进行归一化，并不对形状进行改变，因此要keepdim = True
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

print(20*"=" + "万能的分隔符"+ 20*"=")

features = d_model = 512
eps = 1e-6

x = ff_result

ln = LayerNorm(features, eps)
ln_result = ln(ff_result)
print("ln_result :", ln_result)
print("ln_resutl.size :", ln_result.size())
'''
ln_result : tensor([[[-0.0178,  0.8013,  0.8365,  ...,  2.0148, -0.2267,  0.2413],
         [ 0.2516, -0.3289,  0.1384,  ...,  0.9686, -0.1838, -1.1290],
         [-0.4528,  0.3258,  0.4359,  ...,  2.1818, -0.8507,  0.4520],
         [ 0.4430, -0.1013,  0.3029,  ..., -0.1779,  0.2689, -0.0587]],
        [[ 0.4788,  0.8728,  0.3402,  ..., -0.3265, -1.2746, -0.8108],
         [-0.7220,  1.3556,  1.4107,  ...,  0.5358, -0.1868,  0.9251],
         [ 0.3514,  0.0041,  0.1224,  ...,  0.5421, -1.8704, -0.5819],
         [ 1.1094, -0.8367,  0.5160,  ...,  0.6232,  0.8581, -0.2768]]],
       grad_fn=<AddBackward0>)
ln_resutl.size : torch.Size([2, 4, 512])
'''

# 7.子层连接结构 残差连接
class Sublayerconnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        '''
        :param size: 词嵌入维度
        :param dropout: 防止过拟合
        '''
        super(Sublayerconnection, self).__init__()
        self.size = size
        # 实例化一个归一化层对象
        self.norm = LayerNorm(size)
        # 实例化一个dropout对象
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, sublayer):
        '''将x进行规范化，送入子层函数进行处理，处理的结果进行dropout,最后进行残差连接'''
        return self.norm(x + self.dropout(sublayer(x)))

print(20*"=" + "万能的分隔符"+ 20*"=")

size = d_model = 512
head = 8
x = pe_result
mask = Variable(torch.zeros(8, 4, 4))
self_atten = MultiHeadeAttention(head, d_model)

sublayer = lambda x: self_atten(x, x, x, mask)

sbc = Sublayerconnection(size)
sbc_result = sbc(x, sublayer)

print("sbc_result :", sbc_result)
print("sbc_result.size :", sbc_result.size())
'''
sbc_result : tensor([[[  8.8274,  39.3323,  26.7249,  ...,  -0.2071, -18.8127,   2.5243],
         [ -3.9949, -41.0514,   0.1450,  ...,  14.7426,   0.0000, -11.8379],
         [ 16.0589, -13.3930,  11.8961,  ...,  15.4228,  -2.1990, -10.1740],
         [ 30.5587, -24.7291, -24.7448,  ...,  -6.4760,   5.7037,  35.5374]],
        [[  2.5341,  21.8811, -30.4193,  ...,   0.2469,   0.1950,  20.9492],
         [ 20.0880,  -5.9335,   0.1573,  ...,  -4.3792, -17.3402, -15.6323],
         [  8.1116, -17.6003,  13.5855,  ...,  -9.3609,  14.1211, -10.8686],
         [ 22.5212, -48.2890,  13.1824,  ...,  14.3714,   0.2667,   4.0014]]],
       grad_fn=<AddBackward0>)
sbc_result.size : torch.Size([2, 4, 512])
'''

# 8.编码器层
class EncodeLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncodeLayer, self).__init__()

        self.size = size
        self.self_atten = self_atten
        self.feed_forward = feed_forward

        # 子层连接结构需要两个
        self.sublayer = clones(Sublayerconnection(size, dropout), 2)

    def forward(self, x, mask):
        # x表示上一层传入的张量
        '''x进入第一个子层连接结构，包含多头注意力的子层'''
        '''再让张量经过第二个子层连接结构，包含FF网络'''
        x = self.sublayer[0](x, lambda x: self_atten(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

print(20*"=" + "万能的分隔符"+ 20*"=")

size = d_model = 512
head = 8
d_ff = 64
x = pe_result
dropout = 0.1

self_atten = MultiHeadeAttention(head, d_model)
feed_forward = PositionwiseForward(d_model, d_ff, dropout)
mask = Variable(torch.zeros(8, 4, 4))

el = EncodeLayer(size, self_atten, feed_forward, dropout)

el_result = el(x, mask)
print("el_result :", el_result)
print("el_result.size :", el_result.size())

'''
el_result : tensor([[[-6.8545e+00,  2.8454e+01, -3.1110e+01,  ...,  2.4140e+01,
           3.8591e+01, -5.4986e+00],
         [-6.4183e+00, -2.4148e+01,  2.5211e-01,  ...,  2.0584e-01,
           4.0969e+01, -3.8047e+01],
         [-4.3674e+01, -8.6906e-01, -1.8955e+01,  ..., -6.2761e+00,
          -4.6063e+01, -4.8167e-01],
         [-1.4642e+01, -1.7352e+01,  1.7267e+01,  ..., -3.4092e+01,
          -9.1755e+00, -1.9297e+01]],
        [[-8.5818e+00, -1.2969e+01,  9.6323e+00,  ...,  1.6871e+01,
           4.1137e+00, -2.4861e+01],
         [ 1.3248e+00,  8.5322e-02, -4.4033e+01,  ...,  2.5884e+01,
           4.3539e+01, -6.5779e+00],
         [-1.0634e-02,  8.0901e+00, -3.4857e+01,  ..., -4.0728e+00,
           6.9657e+01, -2.4741e+01],
         [-1.1228e+01, -2.3613e+00,  1.2498e+01,  ...,  4.2794e+00,
          -1.3225e+01,  1.5511e-01]]], grad_fn=<AddBackward0>)
el_result.size : torch.Size([2, 4, 512])
'''

# 9.编码器
class Encoder(nn.Module):
    def __init__(self, layer, N):
        '''layer: 编码器层， N：有几个layer'''
        super(Encoder, self).__init__()
        # 复制N个编码器层
        self.layers = clones(layer, N)
        # 最后用一个归一化层作用在最后一个编码器的输出
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        '''x: 上一层输出的张量，     让x依次经过N个编码器的处理，最后规一化'''
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

print(20*"=" + "万能的分隔符"+ 20*"=")

size = d_model = 512
head = 8
d_ff = 64
dropout = 0.1
x = pe_result

c = copy.deepcopy
self_atten = MultiHeadeAttention(head, d_model)
feed_forward = PositionwiseForward(d_model, d_ff)

layer = EncodeLayer(size, c(self_atten), c(feed_forward), dropout)
N = 8
mask = Variable(torch.zeros(8, 4, 4))

en = Encoder(layer, N)
en_result = en(x, mask)
print("en_result :", en_result)
print("en_result.size :", en_result.size())
'''
en_result : tensor([[[-0.2852, -1.1833,  0.5881,  ..., -2.5690, -0.6298, -2.8499],
         [-1.7299, -1.1440,  0.5023,  ...,  0.1088,  0.5556,  1.4837],
         [-0.4232,  1.9891, -0.3218,  ...,  1.4767,  0.6632,  0.1475],
         [-0.3925, -0.1306, -0.3481,  ..., -0.5113,  0.3718, -1.5519]],
        [[ 0.0751, -1.8860, -0.2580,  ...,  0.7112,  2.3647,  1.6022],
         [ 1.1177, -0.3873,  0.2721,  ...,  0.0565, -1.2874,  0.3400],
         [-0.5698, -0.4667,  1.3811,  ..., -0.2070, -0.0158,  0.0208],
         [-1.2950,  0.8007,  0.4031,  ..., -0.5634, -0.0875, -0.1952]]],
       grad_fn=<AddBackward0>)
en_result.size : torch.Size([2, 4, 512])
'''